from rknn.api import RKNN
import cv2
import time
import numpy as np

if __name__ == '__main__':

    # Create RKNN object
    rknn = RKNN()
   
    # pre-process config
    print('--> config model')
    # 配置模型输入，用于NPU对数据输入的预处理
    # channel_mean_value='0 0 0 255'，那么模型推理时，将会对RGB数据做如下转换
    # (R - 0)/255, (G - 0)/255, (B - 0)/255。推理时，RKNN模型会自动做均值和归一化处理
    # reorder_channel=’0 1 2’用于指定是否调整RBG顺序，设置成0 1 2即按输入的RGB顺序不做调整
    # reorder_channel=’2 1 0’表示交换0和2通道，輸入调整为BGR   
    rknn.config(channel_mean_value='0 0 0 255', reorder_channel='2 1 0')
    print('done')

    # Load tensorflow model
    print('--> Loading model')
    ret = rknn.load_caffe(model='./pose/coco/pose_deploy_linevec.prototxt', proto='caffe',
                            blobs='./pose/coco/pose_iter_440000.caffemodel')
    if ret != 0:
        print('Load model failed!')
        exit(ret)
    print('done')

    # Build model
    print('--> Building model')
        # do_quantization=True进行量化
        # 量化会减小模型的体积和提升运算速度，但是会有精度的丢失
        # 开启预编译，提高load速度
    ret = rknn.build(do_quantization=True, dataset='./dataset.txt', pre_compile=False)
    if ret != 0:
        print('Build model failed!')
        exit(ret)
    print('done')
    # Export rknn model
    print('--> Export RKNN model')
    ret = rknn.export_rknn('./pose_deploy_linevec_pre_compile.rknn')
    if ret != 0:
        print('Export model failed!')
        exit(ret)
    print('done')

    rknn.release()